package com.example.digitclassifier3
import android.content.Context
import android.content.res.AssetManager
import android.graphics.Bitmap
import android.os.FileUtils
import android.util.Log
import com.google.android.gms.tasks.Task
import com.google.android.gms.tasks.TaskCompletionSource
import java.io.FileInputStream
import java.io.IOException
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.channels.FileChannel
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import org.tensorflow.lite.Interpreter

class DigitClassifier(private val context: Context){

    private var interpreter: Interpreter? = null
    var isInitialized = false
    private set

    private val executorService: ExecutorService = Executors.newCachedThreadPool()

    private var inputImageWidth: Int = 0
    private var inputImageHeight: Int = 0
    private var modelInputSize: Int = 0

    fun initialize():Task<Void>{
        //分类器初始化

        val task = TaskCompletionSource<Void>()
        //设置分类器

        executorService.execute{
            try {
                //初始化解释器
                initializeInterpreter()
                task.setResult(null)
            }catch (e:IOException){
                //处理异常
                task.setException(e)
            }
        }

        return task.task
        //返回任务
    }

    @Throws(IOException::class)
    private fun initializeInterpreter(){

        val assetManager = context.assets
        val model = loadModelFile(assetManager,"mnist.tflite")
        //加载模型文件

        val interpreter = Interpreter(model)
        //设置解释器

        val inputShape = interpreter.getInputTensor(0).shape()
        //设置输入格式

        inputImageWidth = inputShape[1]
        inputImageHeight = inputShape[2]
        modelInputSize = FLOAT_TYPE_SIZE * inputImageWidth * inputImageHeight * PIXEL_SIZE
        //设置输入数据的长、宽、大小

        this.interpreter = interpreter
        //将解释器

        isInitialized = true
        //将已初始化的标签设置为真

    }

    @Throws(IOException::class)
    private fun loadModelFile(assetManager: AssetManager,filename:String):ByteBuffer {
        //加载模型文件

        val fileDescriptor = assetManager.openFd(filename)
        //设置文件描述符

        val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
        //设置输入流

        val fileChannel = inputStream.channel
        //设置文件通道

        val startOffset = fileDescriptor.startOffset
        //设置文件起始位置

        val declaredLength = fileDescriptor.declaredLength
        //声明长度

        return fileChannel.map(FileChannel.MapMode.READ_ONLY,startOffset,declaredLength)
        //返回文件通道映射
    }

    private fun classify(bitmap: Bitmap):String {
        check(isInitialized) { "TF Lite Interpreter is not initialized yet." }
        //检查分类器是否已经初始化

        val resizedImage = Bitmap.createScaledBitmap(
            bitmap,
            inputImageWidth,
            inputImageHeight,
            true
        )
        //设置数据格式

        val byteBuffer = convertBitmapToByteBuffer(resizedImage)
        //转化图片数据格式

        val output = Array(1) { FloatArray(OUTPUT_CLASSES_COUNT) }
        //设置结果格式

        interpreter?.run(byteBuffer,output)
        //进行推断

        val result = output[0]
        val maxIndex = result.indices.maxByOrNull { result[it] } ?: -1
        val resultString = "预测结果：%d\n 精度：%2f".format(maxIndex,result[maxIndex])
        //将结果转化为所需要的格式

        return resultString
        //返回结果
    }

    fun classifyAsync(bitmap: Bitmap): Task<String> {
        val task = TaskCompletionSource<String>()
        executorService.execute{
            val result = classify(bitmap)
            task.setResult(result)
        }
        return task.task
    }

    fun close(){
        //设置关闭动作
        executorService.execute{
            interpreter?.close()
            //关闭解释器

            Log.d(TAG,"Closed TFLite interpreter.")
            //记录操作
        }
    }

    private fun convertBitmapToByteBuffer(bitmap: Bitmap):ByteBuffer {
        val byteBuffer = ByteBuffer.allocateDirect(modelInputSize)
        byteBuffer.order(ByteOrder.nativeOrder())
        //设置字节缓冲区

        val pixels = IntArray(inputImageWidth * inputImageHeight)
        bitmap.getPixels(pixels,0,bitmap.width,0,0,bitmap.width,bitmap.height)
        //获取像素数据

        for (pixelValue in pixels){
            val r = (pixelValue shr 16 and 0xFF)
            val g = (pixelValue shr 8 and 0xFF)
            val b = (pixelValue and 0xFF)
            //分别提取红、绿、蓝三通道的数据

            val normalizedPixelValue = (r+g+b) / 3.0f / 255.0f
            //将三通道颜色转化为灰度图并归一化

            byteBuffer.putFloat(normalizedPixelValue)
            //转化数据类型
        }
        return byteBuffer
        //返回字节缓冲区
    }

    companion object {
        private const val TAG = "DigitClassifier"
        private const val FLOAT_TYPE_SIZE = 4
        private const val PIXEL_SIZE = 1
        private const val OUTPUT_CLASSES_COUNT = 10
    }

}
